import gradio as gr
import requests
from PIL import Image
import numpy as np

# FastAPI endpoint URLs
FASTAPI_URL = "http://127.0.0.1:8000/generate-text"
CLEAR_HISTORY_URL = "http://127.0.0.1:8000/clear-history"

def chat_with_model(image: np.ndarray, text: str, conversation_history: list):
    # Convert numpy array image to PIL image format
    pil_image = Image.fromarray(image)

    # Prepare files and data to send to FastAPI
    files = {'image': pil_image}
    data = {'text': text}

    # Send the data to FastAPI for processing
    response = requests.post(FASTAPI_URL, files=files, data=data)
    
    # If the response is successful, return the output text
    if response.status_code == 200:
        output_text = response.json()["output_text"][0]
        
        # Append the assistant's response to the conversation history
        conversation_history.append({
            "role": "assistant",
            "content": [{"type": "text", "text": output_text}],
        })
        
        return conversation_history, output_text
    else:
        return conversation_history, "Error: Unable to get a response from FastAPI."

def clear_conversation_history():
    # Trigger FastAPI to clear the conversation history
    response = requests.post(CLEAR_HISTORY_URL)
    
    if response.status_code == 200:
        return [], "Conversation history cleared."
    else:
        return [], "Error: Failed to clear conversation history."

# Define Gradio interface
iface = gr.Interface(
    fn=chat_with_model,
    inputs=[
        gr.Image(type="numpy", label="Upload an Image"), 
        gr.Textbox(label="Enter your text"),
        gr.State([])  # Maintain conversation history
    ],
    outputs=[
        gr.State([]),  # Output the conversation history
        gr.Textbox(label="Response")  # Display the model's response
    ],
    live=True,
    title="Image and Text Chat with Qwen2.5",
    description="Upload an image and enter text to interact with the model. The model will respond based on both inputs.",
)

# Add the "Clear Conversation History" button
clear_btn = gr.Button("Clear Conversation History")

clear_btn.click(clear_conversation_history, outputs=[iface.inputs[2], iface.outputs[1]])

# Launch the Gradio app
def start_gradio():
    iface.launch()

# When you want to run Gradio, you can call start_gradio()
